package main

import (
	"os"
	"path/filepath"
	"strings"
	"go/ast"
	"go/parser"
	"go/token"
	"sort"
	"io/ioutil"
	"log"
	"fmt"
)

var printArg = "false"

func main() {

	if len(os.Args) < 2 {
		log.Println("没有设置目录")
		return
	}

	dirPath := os.Args[1]

	log.Println("dir path=", dirPath)

	dirWalk := func(path string, info os.FileInfo, err error) error {
		if strings.HasSuffix(path, ".go") && ! strings.HasSuffix(path, "_test.go") {
			modifySourceFile(path)
		}

		return nil
	}

	filepath.Walk(dirPath, dirWalk)

}

func modifySourceFile(path string) {

	log.Println("go file: " + path)

	sourceCode, _ := ioutil.ReadFile(path)

	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, path, sourceCode, parser.ParseComments)

	if err != nil {
		log.Fatal(err)
		return
	}

	insertImportCode := false

	insertCodePosArray := [] int{}
	insertCodeMap := map[int]string{}

	for _, decl := range file.Decls {
		funcDecl, ok := decl.(*ast.FuncDecl)

		if ok {

			if funcDecl.Body == nil {
				continue
			}

			methodName := funcDecl.Name
			argNames := []string{}
			argValues := []string{}

			if funcDecl.Type.Params.List != nil {
				for _, field := range funcDecl.Type.Params.List {
					for _, fieldName := range field.Names {
						if "_" == fieldName.Name {
							continue
						}

						argNames = append(argNames, "\"" + fieldName.Name + "\"")
						argValues = append(argValues, fieldName.Name)
					}
				}
			}

			argNameString := fmt.Sprintf("[]string{%s}", strings.Join(argNames, ","))
			argValueString := fmt.Sprintf("[]interface{}{%s}", strings.Join(argValues, ","))

			insertCodePos := int(funcDecl.Body.Lbrace)
			insertCodePosArray = append(insertCodePosArray, insertCodePos)
			insertCodeMap[insertCodePos] =
				fmt.Sprintf(
					"_gotrace.MethodEntry(\"%s\", %s, %s, %s);", methodName, argNameString, argValueString, printArg) +
					fmt.Sprintf("defer _gotrace.MethodExit(\"%s\");", methodName)

			insertImportCode = true

		}
	}

	if insertImportCode {
		insertCodePos := int(file.Name.End() - 1)
		insertCodePosArray = append(insertCodePosArray, insertCodePos)
		insertCodeMap[insertCodePos] = " ;import _gotrace \"git.oschina.net/jscode/gotrace/trace\""
	}

	sort.Ints(insertCodePosArray)
	for index := len(insertCodePosArray) - 1; index >= 0; index-- {
		insertCodePos := insertCodePosArray[index]
		insertCode := insertCodeMap[insertCodePos]

		sourceCode = insertSlice(sourceCode, insertCodePos, []byte(insertCode))
	}

	ioutil.WriteFile(path, sourceCode, 0)

}

func insertSlice(slice []byte, index int, sub []byte) []byte {

	result := make([]byte, len(slice) + len(sub))

	copy(result, slice[0:index])
	copy(result[index:], sub)
	copy(result[index + len(sub):], slice[index:])

	return result

}


